
import torch
import numpy as np

data = np.loadtxt('data-02-stock_daily.csv', delimiter=',')
data = data[::-1]
from sklearn.preprocessing import MinMaxScaler

data = MinMaxScaler().fit_transform(data)
c = 7
x = []
y = []
for i in range(len(data) - c):
    x_data = data[i:i + c]
    y_data = data[i + c][-1]
    x.append(x_data)
    y.append(y_data)
x = torch.Tensor(x)
y = torch.Tensor(y).reshape(-1, 1)
print()
from sklearn.model_selection import train_test_split

train_x, test_x, train_y, test_y = train_test_split(x, y, shuffle=False)
print(train_x.shape)

input_shape = 5


class Rnn(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.rnn1 = torch.nn.RNN(input_size=x.shape[2], hidden_size=5, batch_first=True)
        self.rnn2 = torch.nn.RNN(input_size=5, hidden_size=10, batch_first=True)
        self.fc = torch.nn.Linear(in_features=10, out_features=1)

    def forward(self, x):
        x, _ = self.rnn1(x)
        x, _ = self.rnn2(x)
        out = self.fc(x[:, -1, :])
        return out


model = Rnn()
loss_fn = torch.nn.MSELoss()
op = torch.optim.Adam(model.parameters(), lr=0.01)
loss_list = []
for i in range(1000):
    op.zero_grad()
    h = model(train_x)
    loss = loss_fn(h, train_y)
    loss.backward()
    loss_list.append(loss)
    op.step()
pre = model(test_x).reshape(-1)
print(pre)
import matplotlib.pyplot as plt

plt.plot(pre.data.numpy(), c='r')
plt.plot(test_y.data.numpy(), c='g')
plt.show()
